import threading
from typing import Any, Dict, List, Optional, Union

import msgspec
import zmq
from msgspec.msgpack import Decoder

from src.index_tree import IndexTree
from src.logging import get_logger
from src.config import get_config


class EventBatch(msgspec.Struct, array_like=True, omit_defaults=True, gc=False):
    ts: float
    events: list[Any]


class KVCacheEvent(
    msgspec.Struct, array_like=True, omit_defaults=True, gc=False, tag=True
):
    """Base class for all KV cache-related events"""


class BlockStored(KVCacheEvent):
    block_hashes: list[int]
    parent_block_hash: Optional[int]
    token_ids: list[int]
    block_size: int
    lora_id: Optional[int]


class BlockRemoved(KVCacheEvent):
    block_hashes: list[int]


class AllBlocksCleared(KVCacheEvent):
    pass


class KVEventBatch(EventBatch):
    events: list[Union[BlockStored, BlockRemoved, AllBlocksCleared]]


def split_store_event(event: BlockStored) -> List[BlockStored]:
    event_size = len(event.block_hashes)
    if event_size == 1:
        return [event]
    res = [None] * event_size
    for i in range(event_size):
        res[i] = BlockStored(block_hashes=[event.block_hashes[i]],
                parent_block_hash=None if i == 0 else event.block_hashes[i - 1],
                token_ids=event.token_ids[i * event.block_size:(i + 1) * event.block_size],
                block_size=event.block_size,
                lora_id=event.lora_id)
    return res


class ZMQWorker(threading.Thread):
    def __init__(self, endpoint: str, trees: List[IndexTree]):
        super().__init__(daemon=True)
        self.endpoint = endpoint
        self.trees = trees  # 该worker专属的tree列表
        self.context = zmq.Context()
        self.running = False
        self.decoder = Decoder(type=KVEventBatch)
        self.lock = threading.Lock()  # 保护index_tree操作
        self.logger = get_logger(f"zmq_worker.{endpoint}")
        
        # 设置线程名称以便日志识别
        self.name = f"ZMQWorker-{endpoint}"
        
    def run(self):
        self.running = True
        self.logger.info("ZMQ Worker starting", endpoint=self.endpoint)
        
        sub = self.context.socket(zmq.SUB)
        sub.connect(self.endpoint)
        sub.setsockopt_string(zmq.SUBSCRIBE, "kv-events")
        
        self.logger.info("ZMQ Worker connected and subscribed", endpoint=self.endpoint)
        
        try:
            while self.running:
                if sub.poll(100):  # 100ms超时
                    _, _, payload = sub.recv_multipart()
                    event_batch = self.decoder.decode(payload)
                    
                    if event_batch.events:
                        self.logger.debug("Processing event batch",
                                        endpoint=self.endpoint,
                                        event_count=len(event_batch.events))
                        self._process_message(event_batch)
                    else:
                        self.logger.warn("Received empty event batch", endpoint=self.endpoint)
        except Exception as e:
            self.logger.error("ZMQ Worker error", endpoint=self.endpoint, error=str(e))
        finally:
            sub.close()
            self.logger.info("ZMQ Worker stopped", endpoint=self.endpoint)
            
    def _process_message(self, event_batch: KVEventBatch):
        try:
            for event in event_batch.events:
                if isinstance(event, BlockStored):
                    self.logger.debug("Processing BlockStored event",
                                    endpoint=self.endpoint,
                                    block_count=len(event.block_hashes),
                                    block_size=event.block_size)
                    self._handle_block_stored(event)
                elif isinstance(event, BlockRemoved):
                    self.logger.debug("Processing BlockRemoved event",
                                    endpoint=self.endpoint,
                                    block_count=len(event.block_hashes))
                    self._handle_block_removed(event)
                elif isinstance(event, AllBlocksCleared):
                    self.logger.debug("Processing AllBlocksCleared event", endpoint=self.endpoint)
                    self._handle_all_cleared()
                else:
                    self.logger.warn("Unknown event type",
                                   endpoint=self.endpoint,
                                   event_type=type(event).__name__)
        except Exception as e:
            self.logger.error("Message processing error",
                            endpoint=self.endpoint,
                            error=str(e))

    def _handle_block_stored(self, event: BlockStored):
        with self.lock:
            single_events = split_store_event(event)
            for e in single_events:
                # build a new tree if it's first block
                if e.parent_block_hash is None:
                    tree = IndexTree(e.block_size)
                    tree.insert_node(e.token_ids, e.parent_block_hash, e.block_hashes[0])
                    self.trees.append(tree)
                    self.logger.info("Created new IndexTree",
                                endpoint=self.endpoint,
                                block_size=e.block_size,
                                tree_count=len(self.trees))
                    continue
                
                inserted = False
                for tree in self.trees:
                    try:
                        tree.insert_node(e.token_ids, e.parent_block_hash, e.block_hashes[0])
                        self.logger.info("insert node to IndexTree ",
                                endpoint=self.endpoint,
                                block_size=e.block_size,
                                tree_count=len(self.trees))
                        inserted = True
                        break
                    except Exception as e:
                        self.logger.error("Failed to insert block into tree",
                                        endpoint=self.endpoint,
                                        error=str(e))
                        continue
                
                if not inserted:
                    self.logger.warn("Failed to insert block into any tree",
                                endpoint=self.endpoint,
                                parent_hash=e.parent_block_hash,
                                curr_hash=e.block_hashes)


    def _handle_block_removed(self, event: BlockRemoved):
        with self.lock:
            removed_count = 0
            for tree in self.trees:
                for block_hash in event.block_hashes:
                    try:
                        # 查找要删除节点的父节点哈希
                        parent_node = tree.find_parent_by_curr_hash(block_hash)
                        if parent_node:
                            tree.delete_node(
                                pre_hash=parent_node.data['curr_hash'],
                                curr_hash=block_hash
                            )
                            removed_count += 1
                        else:
                            # 如果是根节点，特殊处理
                            if tree.root and tree.root.data['curr_hash'] == block_hash:
                                tree.root = None
                                removed_count += 1
                    except Exception as ex:
                        self.logger.warn("Failed to remove block",
                                       endpoint=self.endpoint,
                                       block_hash=block_hash,
                                       error=str(ex))
            
            self.logger.info("Blocks removed",
                           endpoint=self.endpoint,
                           requested_count=len(event.block_hashes),
                           removed_count=removed_count)

    def _handle_all_cleared(self):
        with self.lock:
            if self.trees:
                old_tree_count = len(self.trees)
                block_size = self.trees[0].block_size
                
                # 记录清空前的树状态（如果启用DEBUG）
                if self.trees and self.trees[0].visualizer.config.is_tree_debug_enabled():
                    for i, tree in enumerate(self.trees):
                        if tree.root:  # 只记录非空树
                            debug_log = tree.visualizer.generate_tree_debug_log(tree, "CLEAR")
                            if debug_log:
                                self.logger.debug(f"Tree {i} before CLEAR:\n{debug_log}")
                
                self.trees = [IndexTree(block_size)]
                self.logger.info("All blocks cleared",
                               endpoint=self.endpoint,
                               old_tree_count=old_tree_count,
                               block_size=block_size)
            else:
                self.logger.warn("Attempted to clear empty tree list", endpoint=self.endpoint)


class ZMQManager:
    def __init__(self, trees: Dict[str, List[IndexTree]]):
        self.trees: Dict[str, List[IndexTree]] = trees
        
        # 使用配置管理器获取ZMQ协议
        config_manager = get_config()
        zmq_config = config_manager.zmq
        self.protocol = zmq_config.protocol
        
        self.workers: Dict[str, ZMQWorker] = {}
        self.logger = get_logger("zmq_manager")
        
        # 初始化时创建所有worker
        for endpoint in trees.keys():
            self._create_worker(endpoint)
            
        self.logger.info("ZMQManager initialized",
                        endpoint_count=len(trees),
                        endpoints=list(trees.keys()))
            
    def start(self):
        """启动所有预初始化的worker线程"""
        self.logger.info("Starting ZMQ workers", worker_count=len(self.workers))
        for endpoint, worker in self.workers.items():
            worker.start()
            self.logger.info("Started ZMQ worker", endpoint=endpoint)
            
    def stop(self):
        """停止所有worker线程"""
        self.logger.info("Stopping ZMQ workers", worker_count=len(self.workers))
        for endpoint, worker in self.workers.items():
            worker.running = False
            self.logger.info("Stopping ZMQ worker", endpoint=endpoint)
        
        for endpoint, worker in self.workers.items():
            worker.join()
            self.logger.info("ZMQ worker stopped", endpoint=endpoint)

    def _create_worker(self, endpoint: str):
        """创建指定endpoint的worker"""
        if endpoint not in self.workers:
            worker = ZMQWorker(f"{self.protocol}://{endpoint}", self.trees[endpoint])
            self.workers[endpoint] = worker
            self.logger.info("Created ZMQ worker", endpoint=endpoint)

    def add_worker(self, endpoint: str):
        """添加新的worker"""
        if endpoint not in self.workers:
            self._create_worker(endpoint)
            self.workers[endpoint].start()
            self.logger.info("Added and started ZMQ worker", endpoint=endpoint)
        else:
            self.logger.warn("Worker already exists", endpoint=endpoint)

    def remove_worker(self, endpoint: str):
        """移除指定的worker"""
        if endpoint in self.workers:
            worker = self.workers[endpoint]
            worker.running = False
            worker.join(timeout=5)  # 等待最多5秒
            del self.workers[endpoint]
            self.logger.info("Removed ZMQ worker", endpoint=endpoint)
        else:
            self.logger.warn("Worker not found", endpoint=endpoint)

    def get_tree(self, endpoint: str) -> Optional[IndexTree]:
        """获取指定endpoint的IndexTree"""
        trees = self.trees.get(endpoint)
        if trees:
            return trees[0] if trees else None
        else:
            self.logger.warn("No trees found for endpoint", endpoint=endpoint)
            return None